# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question1/datasets.py
import csv

import scipy.constants
import torch.utils.data

import torch_plus


class EpsDataset(torch.utils.data.Dataset):
    def __init__(self):
        self._data: list[tuple[torch.Tensor, torch.Tensor]] = []
        two = torch_plus.as_tensor(2)
        l_to_omega_factor = two * torch.pi * scipy.constants.c * 1e6 / 1e15
        with open("./data/problem1.csv", newline="") as file:
            reader = csv.reader(file)
            _ = next(reader)
            for row in reader:
                l_ = torch_plus.as_tensor(float(row[0]))
                n = torch_plus.as_tensor(float(row[1]))
                k = torch_plus.as_tensor(float(row[2]))
                self._data.append((
                    l_to_omega_factor / l_,
                    torch.stack([n ** two - k ** two, two * n * k])
                ))

    def __getitem__(self, item):
        return self._data[item]

    def __len__(self):
        return len(self._data)
